import io

from dicom_viewer_be_mock_xiaosai.config import EnumImgTransferMode
from src.dicom_viewer_be_mock_xiaosai.service.api_middleware import logger
from src.dicom_viewer_be_mock_xiaosai.type import EnumRotateSubType
from src.dicom_viewer_be_mock_xiaosai.components import VolumeInfo
from src.dicom_viewer_be_mock_xiaosai.config import img_transfer_mode, EnumImgMode
from src.dicom_viewer_be_mock_xiaosai.layers.BaseImgLayer import BaseImgLayer
from src.dicom_viewer_be_mock_xiaosai.layers.CrossLayer import CrossLayer
from src.dicom_viewer_be_mock_xiaosai.layers.InfoLayer import InfoLayer
from src.dicom_viewer_be_mock_xiaosai.layers.NewMeasureLayer import NewMeasureLayer
from src.dicom_viewer_be_mock_xiaosai.loader.DicomTagLoader import DicomTagLoader
from src.dicom_viewer_be_mock_xiaosai.type import EnumImgQuality, EnumViewType
from src.dicom_viewer_be_mock_xiaosai.utils.util import img_to_jpeg_base64, img_to_png_base64, compose_img


def convert_bytes(image, img_quality):
    with io.BytesIO() as output_bytes:
        if img_quality == EnumImgQuality.high.value:
            image.save(output_bytes, format="PNG", quality=100)
        else:
            image.save(output_bytes, format="JPEG", quality=20)
        img_bytes = output_bytes.getvalue()
    return img_bytes



def img_to_base64(bytes, img_quality):
    if img_quality == EnumImgQuality.high.value:
        return img_to_png_base64(bytes)
    else:
        return img_to_jpeg_base64(bytes)




class Viewer:
    def __init__(self, view_type, dicom_tag_info: DicomTagLoader, volume_info, id_info):
        self.dicom_tag_info = dicom_tag_info
        self.view_type = view_type
        self.volume_info: VolumeInfo = volume_info
        self.base_img_layer = BaseImgLayer(view_type,
                                           dicom_tag_info,
                                           volume_info=volume_info,
                                           id_info= id_info
                                           )  # 基础图像
        self.info_layer = InfoLayer(view_type=self.view_type,
                                    dicom_tag_info=dicom_tag_info, volume_info=self.volume_info)  # 四角信息
        self.cross_layer = CrossLayer(view_type, volume_info=volume_info)
        self.measure_layer = NewMeasureLayer(view_type, dicom_tag_info, volume_info)
        self.mode: EnumImgMode = EnumImgMode.img.value
        self.id_info = id_info




    def toggle_mode(self, new_mode: EnumImgMode):
        self.mode = new_mode

    def get_mix_img(self, img_quality=EnumImgQuality.high.value):
        dicom_data = self.volume_info.get_current_dicom_data(self.view_type)
        return self.get_mix_img_by_dicom_data(dicom_data, img_quality)

    def get_only_info_layer_packed_date(self):
        info_layer_arr = self.info_layer.get_info_layer()
        return self.get_mix_img_by_view_type(None, info_layer_arr, EnumImgQuality.high.value)

    def get_mix_img_by_dicom_data(self, dicom_data, img_quality=EnumImgQuality.high.value):
        base_layer_arr = self.base_img_layer.get_base_img_arr(dicom_data)
        info_layer_arr = self.info_layer.get_info_layer()
        return self.get_mix_img_by_view_type(base_layer_arr, info_layer_arr, img_quality)

    def get_mix_img_by_view_type(self, base_layer_arr, info_layer_arr, img_quality):
        if img_transfer_mode == EnumImgTransferMode.base64.value:
            return self.pack_base64_ans(base_layer_arr, info_layer_arr, img_quality)
        elif img_transfer_mode == EnumImgTransferMode.binary.value:
            return self.pack_binary_ans(base_layer_arr, info_layer_arr, img_quality)
        return None

    def pack_binary_ans(self, base_layer_arr, info_layer_arr, img_quality=EnumImgQuality.high.value):
        match self.view_type:
            case EnumViewType.stack.value:
                if base_layer_arr is not None:
                    other_layer = self.measure_layer.get_layer()
                    mix_rgba_img = compose_img(base_layer_arr, other_layer)

                    img_bytes = convert_bytes(mix_rgba_img, img_quality)

                    current, total = self.get_slice_info()
                    extra_info = self.get_extra_info(self.view_type)
                    return {
                        'img': img_bytes,
                        "info_layer": info_layer_arr,
                        "extra_info": extra_info,
                        "slice_info": {
                            "current": current,
                            "total": total,
                        },
                    }
                else:
                    # 此种情况目前就是提高仅仅是get_position被调用的时候的性能。
                    return {
                        'img': None,
                        "info_layer": info_layer_arr,
                        "slice_info": None
                    }

            case EnumViewType.Sag.value | EnumViewType.Ax.value | EnumViewType.Cor.value:
                if base_layer_arr is not None:
                    center, h_color, v_color, mip = self.cross_layer.get_layer()
                    other_layer = self.measure_layer.get_layer()

                    mix_rgba_img = compose_img(base_layer_arr, other_layer)

                    img_bytes = convert_bytes(mix_rgba_img, img_quality)

                    base64_img = img_to_base64(img_bytes, img_quality)

                    current, total = self.get_slice_info()
                    extra_info = self.get_extra_info(self.view_type)
                    packed_data = {
                        "img": base64_img,
                        "center": {
                            "position": center,
                            "color": {
                                "h_color": h_color,
                                "v_color": v_color,
                            },
                            "mip": mip
                        },
                        "extra_info": extra_info,
                        "info_layer": info_layer_arr,
                        "slice_info": {
                            "current": current,
                            "total": total,
                        },
                    }
                    return packed_data
                else:
                    packed_data = {
                        "img": None,
                        "center": None,
                        "info_layer": info_layer_arr,
                        "slice_info": None
                    }
                    return packed_data
        return None


    def get_extra_info(self, view_type):
        transform_matrix = self.volume_info.view_info[view_type].transform_matrix
        spacing = self.volume_info.get_spacing(view_type)
        return {
            "px_per_mm": transform_matrix[0][0] / float(spacing[0]) / self.volume_info.view_info[view_type].display_to_render_ratio,
            "window_width": self.volume_info.view_info[view_type].window[1],
            "window_center": self.volume_info.view_info[view_type].window[0]
        }

    def pack_base64_ans(self, base_layer_arr, info_layer_arr, img_quality=EnumImgQuality.high.value):
        match self.view_type:
            case EnumViewType.stack.value:
                if base_layer_arr is not None:
                    other_layer = self.measure_layer.get_layer()
                    mix_rgba_img = compose_img(base_layer_arr, other_layer)
                    img_bytes = convert_bytes(mix_rgba_img, img_quality)

                    base64_img = img_to_base64(img_bytes, img_quality)
                    extra_info = self.get_extra_info(self.view_type)
                    current, total = self.get_slice_info()
                    return {
                            "img": base64_img,
                            "info_layer": info_layer_arr,
                            "extra_info": extra_info,
                            "slice_info": {
                                "current": current,
                                "total": total,
                            },
                        }
                else:
                    # 此种情况目前就是提高仅仅是get_position被调用的时候的性能。
                    return {
                        'img': None,
                        "info_layer": info_layer_arr,
                        "slice_info": None
                        }
            case EnumViewType.Sag.value | EnumViewType.Ax.value | EnumViewType.Cor.value:
                if base_layer_arr is not None:
                    center, h_color, v_color, mip = self.cross_layer.get_layer()
                    other_layer = self.measure_layer.get_layer()

                    mix_rgba_img = compose_img(base_layer_arr, other_layer)

                    img_bytes = convert_bytes(mix_rgba_img, img_quality)

                    base64_img = img_to_base64(img_bytes, img_quality)
                    extra_info = self.get_extra_info(self.view_type)
                    current, total = self.get_slice_info()

                    return {
                        "img": base64_img,
                        "center": {
                            "position": center,
                            "color": {
                                "h_color": h_color,
                                "v_color": v_color,
                            },
                            "mip": mip
                        },
                        "extra_info": extra_info,
                        "info_layer": info_layer_arr,
                        "slice_info": {
                            "current": current,
                            "total": total,
                        },
                    }
                else:
                    return {
                        "img": None,
                        "center": None,
                        "info_layer": info_layer_arr,
                        "slice_info": None
                    }
        return None

    def get_slice_info(self):
        current, total = self.volume_info.get_origin_current_total(self.view_type)
        return current + 1, total


    def get_current_mouse_position(self, png_y, png_x):
        origin_dicom_position = self.volume_info.screen_position_to_clipped_voxel_position(self.view_type, png_y,
                                                                                           png_x)
        hu_value = int(self.base_img_layer.get_hu_value_by_position(origin_dicom_position))
        self.info_layer.update_point_value((origin_dicom_position[1] + 1, origin_dicom_position[0] + 1, hu_value))

    def get_recent_img(self, img_quality=EnumImgQuality.low.value):
        dicom_data = self.volume_info.get_current_dicom_data(view_type=self.view_type)
        return self.get_mix_img_by_dicom_data(dicom_data, img_quality)

    def get_recent_without_img(self):
        return self.get_only_info_layer_packed_date()

    def update_flip(self, flip_x, flip_y):
        self.volume_info.transform_viewport_flip(self.view_type, flip_x, flip_y)
        if flip_x:
            self.info_layer.flip_x_direction()
        if flip_y:
            self.info_layer.flip_y_direction()

    def update_rotate(self, rotate_type):
        angle = 90
        clockwise = True
        if rotate_type == EnumRotateSubType.ccw90.value:
            clockwise = False
            self.info_layer.rotate_direction(angle)
        elif rotate_type == EnumRotateSubType.cw90.value:
            self.info_layer.rotate_direction(-angle)
        else:
            logger.error("rotate_type error")
        self.volume_info.transform_viewport_rotate(self.view_type, clockwise, angle)

    def update_window(self, delta_x, delta_y):
        self.volume_info.set_transfer_info_window(self.view_type, delta_x, delta_y)

    def set_window(self, window_width, window_center):
        if window_width == 0 and window_center == 0:
            self.volume_info.init_window(self.view_type)
        else:
            self.volume_info.set_transfer_absolute_window(self.view_type,window_width, window_center)

    def update_zoom(self, zoom_factor):
        self.volume_info.transform_viewport_zoom(self.view_type, zoom_factor)

    def update_pan(self, delta_x, delta_y):
        self.volume_info.transform_viewport_pan(self.view_type, delta_x, delta_y)

    def update_current_slice(self, scroll):
        new_slice_flag = False
        if scroll > 0:
            new_slice_flag = self.volume_info.set_next_slice(self.view_type)
        elif scroll < 0:
            new_slice_flag = self.volume_info.set_prev_slice(self.view_type)
        return new_slice_flag

    def slice_to_index(self, to_index):
        self.volume_info.set_slice_index(self.view_type, to_index - 1)

    def handle_draw(self, sub_op_type, position, action_type):
        self.measure_layer.handle_mouse_op(action_type, position, sub_op_type)

    def reset(self, reset_type):
        match reset_type:
            case 'resetWindow':
                self.reset_window()
            case 'resetMeasure':
                self.reset_measure()
            case 'resetZoom':
                pass
            case 'resetCross':
                pass
            case 'resetAll':
                self.reset_all()
                pass

    def reset_window(self):
        self.volume_info.init_window(self.view_type)

    def reset_measure(self):
        self.measure_layer.clear_measure()

    def reset_pan(self):
        pass

    def reset_zoom(self):
        pass

    def reset_cross(self):
        pass

    def reset_all(self):
        self.reset_window()
        self.reset_measure()
        self.volume_info.init_transform_matrix(self.view_type)

